#include "BRIC.h"
#include "TreeComparator.h"

#include <map>
#include <vector>
#include <cstdio>
#include <algorithm>
#include <opencv/cv.h>

using namespace std;

enum Direction { LEFT, RIGHT };

BRIC::BRIC(int width, int height, int alpha, int beta, int L) {
    this->width = width;
    this->height = height;
    this->alpha = alpha;
    this->beta = beta;
    this->L = L;
}

/*
 TODO: More complex verifications
 */
void BRIC::solve(Node* root) {
    solve(root, VERTICAL);
//    if(root->getArea(alpha, beta) > width * height)
    if(root->getWidth(alpha, beta) > width || root->getHeight(alpha, beta) > height)
        solve(root, HORIZONTAL);
}

void BRIC::calculateNode(CvMat* A, CvMat* B, Node* no, Node* raiz, map<Node*, int>* mapa, Orientation orientacaoSolucao, bool oposto) {
    if(no->isLeaf()) {
        if(oposto) {
            double val = orientacaoSolucao == raiz->getOrientation() ? -1.0 : -(no->getImage()->getAspectRatio());
            cvmSet(A, mapa->at(raiz), mapa->at(no), val);
            val = cvmGet(B, mapa->at(raiz), 0) + 2.0*beta;
            cvmSet(B, mapa->at(raiz), 0, val);
        } else {
            double val = orientacaoSolucao == raiz->getOrientation() ? 1.0 : no->getImage()->getAspectRatio();
            cvmSet(A, mapa->at(raiz), mapa->at(no), val);
            val = cvmGet(B, mapa->at(raiz), 0) - 2.0*beta;
            cvmSet(B, mapa->at(raiz), 0, val);
        }
    } else {
        if(no->getOrientation() == raiz->getOrientation()) {
            calculateNode(A, B, no->getLeft(), raiz, mapa, orientacaoSolucao, oposto);
        } else {
            double val = cvmGet(B, mapa->at(raiz), 0);
            val += oposto ? alpha : -alpha;
            cvmSet(B, mapa->at(raiz), 0, val);
            calculateNode(A, B, no->getLeft(), raiz, mapa, orientacaoSolucao, oposto);
            calculateNode(A, B, no->getRight(), raiz, mapa, orientacaoSolucao, oposto);
        }
    }
}

void BRIC::calculateBissections(CvMat* A, CvMat* B, Node* raiz, map<Node*, int>* mapa, Orientation orientacaoSolucao) {
    if(raiz->getType() == BISS) {
        calculateBissections(A, B, raiz->getLeft(), mapa, orientacaoSolucao);
        calculateBissections(A, B, raiz->getRight(), mapa, orientacaoSolucao);
        calculateNode(A, B, raiz->getLeft(), raiz, mapa, orientacaoSolucao, false);
        calculateNode(A, B, raiz->getRight(), raiz, mapa, orientacaoSolucao, true);
    }
}

/*
 Needed?
 */
int aprox(double d) {
    return ((int)d) + ((d - ((int)d)) > 0.5 ? 1 : 0);
}

void BRIC::applyDimensions(Node* raiz, map<Node*, int>* mapa, CvMat* x, Orientation orientation) {
    if(raiz->isLeaf()) {
        if(orientation == HORIZONTAL) {
            double width = cvmGet(x, mapa->at(raiz), 0);
            double height = width * raiz->getImage()->getAspectRatio();
            raiz->getImage()->setDimensions(aprox(width), aprox(height));
        } else {
            double height = cvmGet(x, mapa->at(raiz), 0);
            double width = height / raiz->getImage()->getAspectRatio();
            raiz->getImage()->setDimensions(aprox(width), aprox(height));
        }
    } else {
        applyDimensions(raiz->getLeft(), mapa, x, orientation);
        applyDimensions(raiz->getRight(), mapa, x, orientation);
    }
}

/*
 TODO:
 -change argument passing: vectors instead of matrixes
 * *how to enumerate terminals ?
 */
void BRIC::solve(Node* root, Orientation orientation) {
    map<Node*, int>* mapa = new map<Node*, int>();
    this->mapear(root, mapa);
    int size = mapa->size()/2 + 1;
    CvMat* A = cvCreateMat(size, size, CV_32FC1);
    CvMat* B = cvCreateMat(size, 1, CV_32FC1);
    CvMat* x = cvCreateMat(size, 1, CV_32FC1);
    Node* temp = new Node(root, orientation, root);
    mapa->insert(pair<Node*, int>(temp, size-1));
    calculateBissections(A, B, root, mapa, orientation);
    cvmSet(B, size-1, 0, (double)(orientation == HORIZONTAL ? width : height));
    calculateNode(A, B, root, temp, mapa, orientation, false);
    cvSolve(A, B, x);
    applyDimensions(root, mapa, x, orientation);
    root->setFather(NULL);
}

/*
 REMOVE
 */
void BRIC::mapear(Node* raiz, map<Node*, int>* mapa) {
    mapa->clear();
    vector<Node*> lista;
    int bissectionNum = 0;
    int leafNum = 0;
    lista.push_back(raiz);
    while(!lista.empty()) {
        if(lista.front()->getType() == BISS) {
            lista.push_back(lista.front()->getLeft());
            lista.push_back(lista.front()->getRight());
            mapa->insert(pair<Node*, int>(lista.front(), bissectionNum++));
        } else {
            mapa->insert(pair<Node*, int>(lista.front(), leafNum++));
        }
        lista.erase(lista.begin());
    }
}


/*
 Returns given node's relative area and aspect ratio, respectively
 */
pair<double, double>* nodeEA(Node* node) {
    pair<double, double>* EA;
    if(node->isLeaf()) {
        EA = new pair<double, double>(node->getImage()->getRelativeArea(), node->getImage()->getAspectRatio());
    } else {
        pair<double, double>* leftEA = nodeEA(node->getLeft());
        pair<double, double>* rightEA = nodeEA(node->getRight());
        double relativeArea;
        double aspectRatio;
        if(node->getOrientation() == VERTICAL) {
            double maxRelativeHeight = max(sqrt(leftEA->first*leftEA->second), sqrt(rightEA->first*rightEA->second));
            double totalRelativeWidth = sqrt(leftEA->first / leftEA->second) + sqrt(rightEA->first / rightEA->second);
            relativeArea = maxRelativeHeight * totalRelativeWidth;
            aspectRatio = maxRelativeHeight / totalRelativeWidth;
        } else {
            double totalRelativeHeight = sqrt(leftEA->first * leftEA->second) + sqrt(rightEA->first * rightEA->second);
            double maxRelativeWidth = max(sqrt(leftEA->first / leftEA->second), sqrt(rightEA->first / rightEA->second));
            relativeArea = maxRelativeWidth * totalRelativeHeight;
            aspectRatio = totalRelativeHeight / maxRelativeWidth;
        }
        EA = new pair<double, double>(relativeArea, aspectRatio);
    }
    return EA;
}

/*
 is it working right?
 */
double BRIC::score2004(Node* node) {
    double aPage = height / (double)width;
    pair<double, double>* EA = nodeEA(node);
    double propE = (node->sumPriorities() / EA->first);
    double propA = (min(EA->second, aPage) / max(EA->second, aPage));
    double score = propE * propA;
    return score;
}

/*
 Is there another method?
 */
void fromRootPath(Node* node, vector<Direction> *path) {
    Node* actual = node;
    while(actual->getFather() != NULL) {
        path->push_back(actual->getFather()->getLeft() == actual ? LEFT : RIGHT);
        actual = actual->getFather();
    }
}

/*
 Re-code ?
 * -solve tree navigation problem (it stinks)
 */
void BRIC::generateTree(Node* root, Node* paramActual, Node* leaf, vector<Node* >* tree_list, Orientation orientation) {
    Node* copy = new Node(root);
    vector<Direction> path;
    fromRootPath(paramActual, &path);
    Node* actual = copy;
    while(!path.empty()) {
        actual = path.back() == LEFT ? actual->getLeft() : actual->getRight();
        path.pop_back();
    }
    Node* father = actual->getFather();
    Node* newNode = new Node(actual, orientation, new Node(leaf));
    if(father) {
        if(father->getLeft() == actual)
            father->setLeft(newNode);
        else
            father->setRight(newNode);
        solve(copy);
        tree_list->push_back(copy);
    } else {
        solve(newNode);
        tree_list->push_back(newNode);
    }
}

void BRIC::generateTrees(Node* root, Node* paramActual, Node* leaf, vector<Node* >* tree_list) {
    generateTree(root, paramActual, leaf, tree_list, HORIZONTAL);
    generateTree(root, paramActual, leaf, tree_list, VERTICAL);
    if(!paramActual->isLeaf()) {
        this->generateTrees(root, paramActual->getLeft(), leaf, tree_list);
        this->generateTrees(root, paramActual->getRight(), leaf, tree_list);
    }
}

double BRIC::individualScore2008(Node* node, double canvas_area, double e_sum) {
    double score = 0.0;
    if(node->isLeaf()) {
        double e_proportion = node->getImage()->getRelativeArea() / e_sum;
        double ideal_area = canvas_area * e_proportion;
        double areaProportion = node->getImage()->getArea() / ideal_area;
        score = min(areaProportion, 1.0);
        if(score < 0.25)
            score = 0.0;
    } else {
        score = individualScore2008(node->getLeft(), canvas_area, e_sum) + individualScore2008(node->getRight(), canvas_area, e_sum);
    }
    return score;
}

/*
 Re-code:
 * -solve solve()'s collateral effect
 * -find a working score calculation (relation between terminal areas considering the mathematical solution)
 */
double BRIC::score2008(Node* root) {
    double sum_e = root->sumPriorities();
    double canvas_area = width * height;
    double score = individualScore2008(root, canvas_area, sum_e);
    return score;
}

void getLeafs(Node* root, vector<Node*>* leafs) {
    if(root->isLeaf())
        leafs->push_back(root);
    else {
        getLeafs(root->getLeft(), leafs);
        getLeafs(root->getRight(), leafs);
    }
}

double BRIC::score2012(Node* root) {
    double area = width * height;
    double node_a = root->getHeight(alpha, beta) / root->getWidth(alpha, beta);
    double canvas_a = height / width;
    double AR = min(node_a, canvas_a) / max(node_a, canvas_a);
    vector<Node*> leafs;
    getLeafs(root, &leafs);
    double e_sum = root->sumPriorities();
    double score = 0.0;
    for(int i = 0; i < leafs.size(); i++) {
        double ideal_area = (leafs.at(i)->getImage()->getRelativeArea() / e_sum) * area;
        score += min(ideal_area, leafs.at(i)->getImage()->getArea()) / max(ideal_area, leafs.at(i)->getImage()->getArea());
    }
    return score; // + AR;
}

double BRIC::score(Node* root) {
//    return score2004(root);
//    return score2008(root);
    return score2012(root);
}

void print_tree(Node* root) {
    if(root) {
        if(root->isLeaf()) {
            printf("[%d *%d]", root->getImage(), root->getFather());
        } else {
            printf("( ");
            print_tree(root->getLeft());
            printf(" %s *%d ", root->getOrientation() == HORIZONTAL ? "H" : "V", root->getFather());
            print_tree(root->getRight());
            printf(" )");
        }
    }
}

/*
 Break into sub problems
 */
Node* BRIC::createTree(vector<Image*>* images) {
    vector<Node*> leaf_list;
    for(int i = 0; i < images->size(); i++) {
        leaf_list.push_back(new Node(images->at(i)));
    }
    vector<Node*> best_trees;
    //treat case
    best_trees.push_back(leaf_list.at(0));
    for(int i = 1; i < leaf_list.size(); i++) {
        vector<Node*> trees;
        for(int j = 0; j < best_trees.size(); j++)
            this->generateTrees(best_trees.at(j), best_trees.at(j), leaf_list.at(i), &trees);
        //capture L or trees.size() different scores
        sort(trees.begin(), trees.begin() + trees.size(), TreeComparator(this));
        best_trees.clear(); //free pointers
        best_trees.push_back(trees.back());
        trees.pop_back();
        double lastScore = score(best_trees.back());
        for(int j = 0; j < (int)trees.size() && (int)best_trees.size() < L; j++) {
            double sc = score(trees.back());
            if(sc != lastScore) {
                best_trees.push_back(trees.back());
                lastScore = sc;
            }
            trees.pop_back();
        }
    }
    //find max
    Node* max = best_trees.at(0);
    for(int a = 1; a < best_trees.size(); a++) {
        if(score(max) < score(best_trees.at(a)))
            max = best_trees.at(a);
    }
    return max;
}
